package com.yml.rust.imageproc;

import android.graphics.Bitmap;
import android.graphics.BitmapFactory;
import android.graphics.Color;
import android.os.Bundle;
import android.os.Handler;
import android.os.Message;
import android.os.PersistableBundle;
import android.util.Base64;
import android.util.Log;
import android.widget.ImageView;
import android.widget.TextView;

import androidx.annotation.NonNull;
import androidx.annotation.Nullable;
import androidx.appcompat.app.AppCompatActivity;

import com.gemalto.jp2.JP2Decoder;

import java.io.BufferedWriter;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.FileWriter;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

public class MainActivity extends AppCompatActivity {

    public static String imageBase64 = "";

    public static final int HANDLER_MSG_1 = 1;

    public static final String TAG = "aiclahe";
    private ImageView originImg, aiImg;
    private TextView testText;

    private Handler handler = new Handler() {
        @Override
        public void handleMessage(@NonNull Message msg) {
            super.handleMessage(msg);
            if (msg.what == HANDLER_MSG_1) {
                testText.setText(msg.obj.toString());
            }
        }
    };

    @Override
    public void onCreate(@Nullable Bundle savedInstanceState, @Nullable PersistableBundle persistentState) {
        super.onCreate(savedInstanceState, persistentState);
    }

    @Override
    protected void onCreate(@Nullable Bundle savedInstanceState) {
        super.onCreate(savedInstanceState);
        setContentView(R.layout.activity_main);
        originImg = findViewById(R.id.origin_img);
        aiImg = findViewById(R.id.ai_img);
        testText = findViewById(R.id.test_text);

        String output = HelloWorld.hello("josh");
        Log.d("HelloWorld", "HelloWorld.hello: " + output);

//        Bitmap originBmp = BitmapFactory.decodeResource(getResources(), R.drawable.finger1);
        String base64Data = imageBase64.substring("data:image/jp2;base64,".length());
        byte[] imageBytes = Base64.decode(base64Data, 0);
        Log.d("HelloWorld", "imageBytes size: " + imageBytes.length);
        Log.d("HelloWorld", "imageBytes width: " + imageBytes[0] + ", height: " + imageBytes[1] + ", alpha: " + imageBytes[2]);
        if (android.os.Build.VERSION.SDK_INT >= android.os.Build.VERSION_CODES.O) {
            Color firstColor = Color.valueOf(imageBytes[3]);
            Log.d("HelloWorld", "imageBytes first color argb: " + firstColor.alpha()
                    + "-" + firstColor.red() + "-" + firstColor.green() + "-" + firstColor.blue());
        }
        Bitmap originBmp = new JP2Decoder(imageBytes).decode();
        originImg.setImageBitmap(originBmp);

//        Bitmap aiBmp = clahe(originBmp, 8, 48);
//        aiImg.setImageBitmap(aiBmp);

        // 测试rust做图片ai增强
        new Thread(new Runnable() {
            @Override
            public void run() {
                String result = ImageProcLib.testByteArrTrans(new byte[]{-128, 127});
                Message msg = new Message();
                msg.what = HANDLER_MSG_1;
                msg.obj  = result;
                handler.sendMessage(msg);
            }
        }).start();


        byte[] raw = convertBitmapToByte(originBmp);
        byte[] aiImageArr = ImageProcLib.clahe(raw, originBmp.getWidth(), originBmp.getHeight(), 8, 48);
        Log.d("ImageProcLib", "aiImageArr length: " + aiImageArr.length);
//        Log.d("ImageProcLib", "aiImageArr: " + Arrays.toString(aiImageArr));
        String aiImageBase64 = Base64.encodeToString(aiImageArr, 0);
        String fileName = "ai.txt";
        File externalStorageDir = getExternalCacheDir();
        File file = new File(externalStorageDir, fileName);
        try {
            file.createNewFile();
            FileWriter writer = new FileWriter(file);
            BufferedWriter bufferedWriter = new BufferedWriter(writer);
            bufferedWriter.write(aiImageBase64); // 写入字符串内容
            bufferedWriter.close();
        } catch (IOException e) {
            e.printStackTrace();
        }

        Bitmap aiBmp = convertToBitmap(aiImageArr);
        aiImg.setImageBitmap(aiBmp);
    }

    /**
     * @param bitmap
     * @param block_number 默认是8
     * @param limit        默认是48
     */
    public Bitmap clahe(Bitmap bitmap, int block_number, int limit) {
        int[] image = convert(bitmap);
        Log.d(TAG, "image length = " + image.length);

        int width = bitmap.getWidth(); // 获取图片宽度
        int height = bitmap.getHeight(); // 获取图片高度
        int block_height = height / block_number;
        Log.d(TAG, "block_height = " + block_height);
        int block_width = width / block_number;
        Log.d(TAG, "block_width = " + block_width);

        // 初始化累积分布函数
        int usize = ((block_number * block_number) << 8);
        List<Float> all_lut = new ArrayList<>(usize); // 容量usize，元素都是0的集合
        for (int i = 0; i < usize; i++) {
            all_lut.add(0.0f);
        }
        Log.d(TAG, "all_lut size = " + all_lut.size());
        int block_size = block_width * block_height;
        int row_incr = (block_height * width) << 2;
        int col_incr = block_width << 2;
        Log.d(TAG, "row_incr = " + row_incr + ", col_incr = " + col_incr);
        int row_start_index = 0;

        // 求得各个块的累积分布函数
        for (int i = 0; i < block_number; i++) {
            int block_start_index = row_start_index;
            for (int j = 0; j < block_number; j++) {
                List<Integer> block = new ArrayList<>(block_size); // 容量block_size，元素都是0的集合
                for (int k = 0; k < block_size; k++) {
                    block.add(0);
                }
                int hist_size = ((i * block_number + j) << 8);
                for (int k = 0; k < block_height; k++) {
                    int start_index = block_start_index + ((k * width) << 2);
                    for (int l = 0; l < block_width; l++) {
                        int block_index = (block_width * k + l);
                        block.set(block_index, image[start_index]); // lock 看着只是对透明度或者是红色做了统计
                        start_index += 4;
                    }
                }
                List<Integer> hist_block = hist_cal(block, block_size);
                equalize_histogram(block_size, hist_block, limit, all_lut, hist_size);
                block_start_index += col_incr;
            }
            row_start_index += row_incr;
        }

        // 插值
        int half_width = block_width >> 1;
        int half_height = block_height >> 1;
        Log.d(TAG, "half_width = " + half_width + ", half_height = " + half_height);
        for (int m = 0; m < height; m++) {
            for (int n = 0; n < width; n++) {
                // four coners
                // 在角上的像素点为累积分布函数得到的值
                int index = (m * width + n) << 2;
                int new_value;
                int current_value = image[index];
                if (m <= half_height && n <= half_width) {
                    int block_m = 0;
                    int block_n = 0;
                    int all_lut_index = (((block_m * block_number + block_n) << 8) + current_value);
                    new_value = all_lut.get(all_lut_index).intValue();
                } else if (m <= half_height && (n >= width - half_width)) {
                    int block_m = 0;
                    int block_n = block_number - 1;
                    int all_lut_index = (((block_m * block_number + block_n) << 8) + current_value);
                    new_value = all_lut.get(all_lut_index).intValue();
                } else if (n <= half_width && m >= height - half_height) {
                    int block_m = block_number - 1;
                    int block_n = 0;
                    int all_lut_index = (((block_m * block_number + block_n) << 8) + current_value);
                    new_value = all_lut.get(all_lut_index).intValue();
                } else if (m >= height - half_height && n >= width - half_width) {
                    int block_m = block_number - 1;
                    int block_n = block_number - 1;
                    int all_lut_index = (((block_m * block_number + block_n) << 8) + current_value);
                    new_value = all_lut.get(all_lut_index).intValue();
                }
                // four edges except coners
                // 四个边采取线性插值
                else if (n <= half_width) {
                    int block_m1 = clamp((m - half_height) / block_height, 0, 6);
                    int block_n1 = 0;
                    int block_m2 = block_m1 + 1;
                    int block_n2 = block_n1;
                    float u = ((float) m - (float) ((block_m1 * block_height) + half_height)) / (float) block_height;
                    float v = 1 - u;
                    int all_lut_indexV = (((block_m1 * block_number + block_n1) << 8) + current_value);
                    int all_lut_indexU = (((block_m2 * block_number + block_n2) << 8) + current_value);
                    new_value = (int) (v * all_lut.get(all_lut_indexV) + u * all_lut.get(all_lut_indexU));
                } else if (m <= half_height) {
                    int block_m1 = 0;
                    int block_n1 = clamp((n - half_width) / block_width, 0, 6);
                    int block_m2 = block_m1;
                    int block_n2 = block_n1 + 1;
                    float u = ((float) n - (float) (block_n1 * block_width + half_width)) / (float) block_width;
                    float v = 1.0f - u;
                    int all_lut_indexV = (((block_m1 * block_number + block_n1) << 8) + current_value);
                    int all_lut_indexU = (((block_m2 * block_number + block_n2) << 8) + current_value);
                    new_value = (int) (v * all_lut.get(all_lut_indexV) + u * all_lut.get(all_lut_indexU));
                } else if (m >= height - half_height) {
                    int block_m1 = block_number - 1;
                    int block_n1 = clamp((n - half_width) / block_width, 0, 6);
                    int block_m2 = block_m1;
                    int block_n2 = block_n1 + 1;
                    float u = ((float) n - (float) (block_n1 * block_width + half_width)) / (float) block_width;
                    float v = 1.0f - u;
                    int all_lut_indexV = (((block_m1 * block_number + block_n1) << 8) + current_value);
                    int all_lut_indexU = (((block_m2 * block_number + block_n2) << 8) + current_value);
                    new_value = (int) (v * all_lut.get(all_lut_indexV) + u * all_lut.get(all_lut_indexU));
                } else if (n >= width - half_width) {
                    int block_m1 = clamp((m - half_height) / block_height, 0, 6);
                    int block_n1 = block_number - 1;
                    int block_m2 = block_m1 + 1;
                    int block_n2 = block_n1;
                    float u = ((float) m - (float) (block_m1 * block_height + half_height)) / (float) block_height;
                    float v = 1.0f - u;
                    int all_lut_indexV = (((block_m1 * block_number + block_n1) << 8) + current_value);
                    int all_lut_indexU = (((block_m2 * block_number + block_n2) << 8) + current_value);
                    new_value = (int) (v * all_lut.get(all_lut_indexV) + u * all_lut.get(all_lut_indexU));
                }
                // 中间区域采取双线性插值
                else {
                    int block_m1 = clamp((m - half_height) / block_height, 0, 6);
                    int block_n1 = clamp((n - half_width) / block_width, 0, 6);
                    int i1 = (((block_m1 * block_number + block_n1) << 8) + current_value);
                    int i2 = i1 + 256;
                    int i3 = i1 + (block_number << 8);
                    int i4 = i3 + 256;
                    float u = (float) (m - ((block_m1 * block_height) + (half_height))) / (float) block_height;
                    float v = (float) (n - ((block_n1 * block_width) + (half_width))) / (float) block_width;
                    float uv = u * v;
                    new_value = (int) ((float) (1.0 - u - v + uv) * all_lut.get(i1) + (float) (v - uv) * all_lut.get(i2) + (float) (u - uv) * all_lut.get(i3) + (float) uv * all_lut.get(i4));
                }
                image[index] = (byte) new_value;
                image[index + 1] = (byte) new_value;
                image[index + 2] = (byte) new_value;
            }
        }

//        return BitmapFactory.decodeByteArray(image, 0, image.length);
        return convertToBitmap(image);
    }

    /**
     * 计算所有块的直方图，就是统计透明度或者红色值[0-255]某一个值的个数
     *
     * @param block  块的集合
     * @param length 块里面像素点的个数，每个像素点由argb四个[0-255]的值构成
     * @return 直方图结果集
     */
    public List<Integer> hist_cal(List<Integer> block, int length) {
        List<Integer> histogram = new ArrayList<>(256); // 容量256，元素都是0的集合
        for (int i = 0; i < 256; i++) {
            histogram.add(0);
        }
        for (int i = 0; i < length; i++) {
            int pixel_value = block.get(i);
            histogram.set(pixel_value, histogram.get(pixel_value) + 1);
        }
        return histogram;
    }

    /**
     * 均衡直方图
     *
     * @param pixcount 块的像素个数
     * @param histogram 直方图集合
     * @param limit 限制
     * @param lut 要填充的lut
     * @param offset 偏移量，由移动到哪个块来计算
     */
    public void equalize_histogram(int pixcount, List<Integer> histogram, int limit, List<Float> lut, int offset) {
        int steal = 0;
        for (int k = 0; k < 256; k++) {
            if (histogram.get(k) > limit) { // 某个亮度值超过限制48个，超过的数量累计到 steal 中，这个亮度的数量设为 limit
                steal += histogram.get(k) - limit;
                histogram.set(k, limit); // 将超过限制的直方图数值都设置为 limit
            }
        }
        int bonus = steal >> 8;
        int res = steal - (bonus << 8);
        Log.d(TAG, "steal = " + steal + ", bonus = " + bonus + ", res = " + res + ", pixcount = " + pixcount);
        if (res > 0) {
            int res_inv = 256 / res;
            int res_step = Math.max(res_inv, 1);
            // 对剩余res循环，并按照res_step间隔对直方图数组进行递增。
            for (int k = 0; k < res; k++) {
                int res_index = k * res_step;
                histogram.set(res_index, histogram.get(res_index) + 1);
            }
        }

        // TODO 优化1，不需要创建这个集合
//        List<Float> pro_den = new ArrayList<>(256); // 容量256，元素都是0的集合
//        for (int i = 0; i < 256; i++) {
//            pro_den.add(0.0f);
//        }
        float sum = 0.0f;
        for (int k = 0; k < 256; k++) {
            histogram.set(k, histogram.get(k) + bonus);
//            pro_den.set(k, (float) histogram.get(k) / (float) pixcount);
//            sum += pro_den.get(k);
            sum += ((float) histogram.get(k) / (float) pixcount);
//            Log.d(TAG, "----- " + histogram.get(k) + ", sum = " + sum);
            lut.set(offset + k, sum * 255.0f);
        }
    }

    /**
     * 当一个数超出限定的范围，返回范围的边界值
     */
    public int clamp(int self, int min, int max) {
        if (self < min) {
            return min;
        }
        return Math.min(self, max);
    }


    public static byte[] convertBitmapToByte(Bitmap bitmap) {
        // 创建与bitmap相同大小的int数组来存放像素值
        int width = bitmap.getWidth();
        Log.d(TAG, "bitmap width = " + width);
        int height = bitmap.getHeight();
        Log.d(TAG, "bitmap height = " + height);
//        int size = width * height * 4;
        int size = bitmap.getByteCount(); // 和 width*height*4 的大小相同
        Log.d(TAG, "bitmap byte size = " + size);


        Log.d("ImageProcLib", "origin image argb pixels size: " + size);
        ByteBuffer buffer = ByteBuffer.allocate(size);
        bitmap.copyPixelsToBuffer(buffer);
        byte[] bitmapByteArr = buffer.array();
        Log.d(TAG, "rgba = " + bitmapByteArr[0] + ", " + bitmapByteArr[1] + ", " + bitmapByteArr[2] + ", " + bitmapByteArr[3]);
        Log.d(TAG, "rgba = " + bitmapByteArr[size - 4] + ", " + bitmapByteArr[size - 3] + ", " + bitmapByteArr[size - 2] + ", " + bitmapByteArr[size - 1]);

        return bitmapByteArr;
    }

    /**
     * 将图片转为 int[]，可能是 rgba 的顺序，可能是 argb 的顺序。
     */
    public static int[] convert(Bitmap bitmap) {
        // 创建与bitmap相同大小的int数组来存放像素值
        int width = bitmap.getWidth();
        Log.d(TAG, "bitmap width = " + width);
        int height = bitmap.getHeight();
        Log.d(TAG, "bitmap height = " + height);
        int[] pixels = new int[width * height];
        Log.d(TAG, "pixels size = " + pixels.length);

        // 使用getPixels()函数获取图片每个像素点的颜色信息并保存到pixels数组中
        bitmap.getPixels(pixels, 0, width, 0, 0, width, height);
        Log.d("ImageProcLib", "origin image pixels: " + Arrays.toString(pixels));

        if (android.os.Build.VERSION.SDK_INT >= android.os.Build.VERSION_CODES.O) {
            Color argbColor = Color.valueOf(pixels[0]);
            Log.d(TAG, "color1 a = " + argbColor.alpha() + ", r = " + argbColor.red() + ", g = " + argbColor.green() + ", b = " + argbColor.blue());
        }

        int size = width * height * 4;
        pixels = new int[size];
        ByteBuffer buffer = ByteBuffer.allocateDirect(size);
        bitmap.copyPixelsToBuffer(buffer);
//        int flag = 0;
        for (int i = 0; i < size; i++) {
            buffer.position(i);
            byte curr = buffer.get();
            pixels[i] = curr;

//            if (curr < 0) {
//                pixels[i] = 256 + curr;
//            } else {
//                pixels[i] = curr;
//            }

            // 现在是rgba的顺序，需要调整为argb
//            flag++;
//            if (flag == 4) {
//                int temp = pixels[i];
//                pixels[i] = pixels[i - 1];
//                pixels[i - 1] = pixels[i - 2];
//                pixels[i - 2] = pixels[i - 3];
//                pixels[i - 3] = temp;
//                flag = 0;
//            }

        }

        Log.d(TAG, "color1 a = " + pixels[0] + ", r = " + pixels[1] + ", g = " + pixels[2] + ", b = " + pixels[3]);

        return pixels;
    }


    public static Bitmap convertToBitmap(int[] rgbArray) {
        int le = rgbArray.length / 4;
        Log.d(TAG, "convertToBitmap le = " + le);

        int width = (int) Math.sqrt(le); // 计算图像宽高
        Log.d(TAG, "convertToBitmap width = " + width);
        if ((width * width != le)) {
            throw new IllegalArgumentException("RGB array length must be a square number");
        }

        Bitmap bitmap = Bitmap.createBitmap(width, width, Bitmap.Config.ARGB_8888);
        for (int i = 0; i < le; i++) {
//            int red = rgbArray[i * 4 + 1];
//            int green = rgbArray[i * 4 + 2];
//            int blue = rgbArray[i * 4 + 3];
//            int alpha = rgbArray[i * 4];

            int red = rgbArray[i * 4];
            int green = rgbArray[i * 4 + 1];
            int blue = rgbArray[i * 4 + 2];
            int alpha = rgbArray[i * 4 + 3];

            int x = i % width;
            int y = i / width;

            bitmap.setPixel(x, y, Color.argb(alpha, red, green, blue));
        }

        return bitmap;
    }


    public  Bitmap convertToBitmap(byte[] rgbArray) {
        int le = rgbArray.length / 4;
        Log.d(TAG, "convertToBitmap rgba byte array length = " + le);

        int width = (int) Math.sqrt(le); // 计算图像宽高
        Log.d(TAG, "convertToBitmap width = " + width);
        if ((width * width != le)) {
            throw new IllegalArgumentException("RGB array length must be a square number");
        }

        Bitmap bitmap = Bitmap.createBitmap(width, width, Bitmap.Config.ARGB_8888);
        for (int i = 0; i < le; i++) {
//            int red = rgbArray[i * 4 + 1];
//            int green = rgbArray[i * 4 + 2];
//            int blue = rgbArray[i * 4 + 3];
//            int alpha = rgbArray[i * 4];

            int red = rgbArray[i * 4];
            int green = rgbArray[i * 4 + 1];
            int blue = rgbArray[i * 4 + 2];
            int alpha = rgbArray[i * 4 + 3];
            if (red < 0) red = red + 256;
            if (green < 0) green = green + 256;
            if (blue < 0) blue = blue + 256;
            if (alpha < 0) alpha = alpha + 256;

            if (i == 0 || i == (le - 4)) {
                Log.d(TAG, "result rgba = " + red + ", " + green + ", " + blue + ", " + alpha);
            }

            int x = i % width;
            int y = i / width;

            bitmap.setPixel(x, y, Color.argb(alpha, red, green, blue));
        }

        return bitmap;
    }

}
